import math

initial_lr = 0.000000001  # 初始学习率
max_lr = 0.00005  # 最大学习率（在warm-up结束时达到）
warmup_epochs = 8
weight_decay = 0.0001
epoch_num = 100  # 训练总轮数


def lr_lambda(epoch):
    if epoch < warmup_epochs:
        # warm-up阶段：线性上升学习率
        x = (max_lr / initial_lr) * (epoch + 1) / warmup_epochs
        return x
    else:
        # 延长的余弦退火阶段，直到训练结束
        adjusted_epoch = epoch - warmup_epochs
        adjusted_total_epochs = 2*epoch_num - warmup_epochs
        progress = adjusted_epoch / adjusted_total_epochs
        # 衰减阶段：使用余弦退火策略
        x = max_lr / initial_lr * 0.5 * (1 + math.cos(math.pi * progress))
        return x


if __name__ == '__main__':
    for epoch in range(epoch_num):
        lr = initial_lr * lr_lambda(epoch)
        print(f"Epoch {epoch + 1}: lr={lr:.6f}")
